package com.mojo.jedis.config;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.connection.ReturnType;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

@Component
public class RateLimiterInterceptor implements HandlerInterceptor {

    private static final String RATE_LIMITER_KEY_PREFIX = "rate_limiter:";
    private static final long DEFAULT_LIMIT = 5; // 默认限流数量
    private static final long DEFAULT_INTERVAL = 5; // 默认限流时间间隔（秒）

    @Autowired
    private StringRedisTemplate redisTemplate;
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String uri = request.getRequestURI();
        String rateLimiterKey = RATE_LIMITER_KEY_PREFIX + uri;
        String limitStr = redisTemplate.opsForValue().get(rateLimiterKey);
        long limit = limitStr != null ? Long.parseLong(limitStr) : DEFAULT_LIMIT;

        // 使用Lua脚本原子性地执行限流逻辑
        String luaScript = "local current = redis.call('incr', KEYS[1])\n"
                + "if tonumber(current) == 1 then\n"
                + "    redis.call('expire', KEYS[1], ARGV[1])\n"
                + "end\n"
                + "return tonumber(current) <= tonumber(ARGV[2])";
        Boolean result = redisTemplate.execute((RedisCallback<Boolean>) connection -> {
            return connection.eval(luaScript.getBytes(), ReturnType.BOOLEAN, 1, rateLimiterKey.getBytes(),
                    String.valueOf(DEFAULT_INTERVAL).getBytes(), String.valueOf(limit).getBytes());
        });
        if (result != null && result) {
            return true;
        } else {
            response.getWriter().write("请求过于频繁");
            return false;
        }
    }
}
